import sys
import threading
import time
import zmq
from zhelpers import zpipe


def client_task(ctx, pipe):
    client = ctx.socket(zmq.DEALER)
    client.identity = b'C'
    client.connect("tcp://localhost:5555")

    print("Setting up test...")
    time.sleep(0.1)

    print("Synchronous round-trip test...")
    start = time.time()
    requests = 10000
    for r in range(requests):
        # print('client send hello %s' % r)
        client.send(b"hello")
        msg = client.recv()
        # print('client recv %s' % msg)
    print(" %d calls/second" % (requests / (time.time()-start)))

    print("Asynchronous round-trip test...")
    start = time.time()
    for r in range(requests):
        client.send(b"hello")
    for r in range(requests):
        client.recv()
    print(" %d calls/second" % (requests / (time.time()-start)))

    pipe.send(b"done")


def worker_task():
    ctx = zmq.Context()
    worker = ctx.socket(zmq.DEALER)
    worker.identity = b'W'
    worker.connect("tcp://localhost:5556")
    # worker.send(b'ready')
    while True:
        msg = worker.recv_multipart()
        # print('worker recv %s' % msg)
        # print('worker send %s' % msg)
        worker.send_multipart(msg)
    ctx.destroy(0)


def broker_task():
    ctx = zmq.Context()
    frontend = ctx.socket(zmq.ROUTER)
    backend = ctx.socket(zmq.ROUTER)
    frontend.bind("tcp://*:5555")
    backend.bind("tcp://*:5556")

    poller = zmq.Poller()
    poller.register(backend, zmq.POLLIN)
    poller.register(frontend, zmq.POLLIN)

    while True:
        try:
            items = dict(poller.poll())
        except KeyboardInterrupt:
            break
        if frontend in items:
            msg = frontend.recv_multipart()
            # print('frontend recv %s ' % msg)
            msg[0] = b'W'
            # print('backend send %s' % msg)
            backend.send_multipart(msg)
        if backend in items:
            msg = backend.recv_multipart()
            # print('backend recv %s' % msg)
            msg[0] = b'C'
            # print('frontend send %s' % msg)
            frontend.send_multipart(msg)


def main():
    ctx = zmq.Context()
    client, pipe = zpipe(ctx)

    client_thread = threading.Thread(target=client_task, args=(ctx, pipe))
    worker_thread = threading.Thread(target=worker_task)
    worker_thread.daemon = True
    broker_thread = threading.Thread(target=broker_task)
    broker_thread.daemon = True

    broker_thread.start()
    worker_thread.start()
    time.sleep(1)
    client_thread.start()

    client.recv()


if __name__ == "__main__":
    main()
